from node2vec.src.model import model
from node2vec.src.trainer.parser import get_parser_args
import networkx as nx
from gensim.models import Word2Vec


def read_graph(args):
    """
    Reads the input network in networkx.
    """
    if args.weighted:
        G = nx.read_edgelist(args.input, nodetype=int, data=(('weight', float),), create_using=nx.DiGraph())
    else:
        G = nx.read_edgelist(args.input, nodetype=int, create_using=nx.DiGraph())
        for edge in G.edges():
            G[edge[0]][edge[1]]['weight'] = 1

    if not args.directed:
        G = G.to_undirected()

    return G


def learn_embeddings(args, walks):
    """
    Learn embeddings by optimizing the Skipgram objective using SGD.
    """
    walks = [list(map(str, walk)) for walk in walks]
    model = Word2Vec(walks, size=args.dimensions, window=args.window_size, min_count=0, sg=1, workers=args.workers,
                     iter=args.iter)
    model.wv.save_word2vec_format(args.output)


def train(args):
    """
    Pipeline for representational learning for all nodes in a graph.
    """
    nx_G = read_graph(args)
    random_walker = model.RandomWalker(nx_G, args.directed, args.p, args.q, 'node2vec_kk')
    walks = random_walker.simulate_walks(args.num_walks, args.walk_length)
    learn_embeddings(args, walks)


def main():
    args = get_parser_args()
    print('training start!')
    train(args)
    # test_auto(args)
    print('training finished!')


if __name__ == '__main__':
    main()
